package com.xgk.boot.framework.common.util.threadpool;

import com.xgk.boot.framework.common.context.ElmOrderThreadContext;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;

public class ElmThreadPoolExecutor extends ThreadPoolExecutor {

    public ElmThreadPoolExecutor(int corePoolSize, int maxPoolSize, long keepAliveTime, TimeUnit unit,
                                 BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maxPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    public ElmThreadPoolExecutor(int corePoolSize, int maxPoolSize, long keepAliveTime, TimeUnit unit,
                                 BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maxPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    @Override
    public void execute(Runnable command) {
        // 包装原始任务
        super.execute(wrap(command));
    }

    public Runnable wrap(Runnable runnable) {
        Map<String, Object> context = new HashMap<>(ElmOrderThreadContext.getAll());
        return () -> {
            try {
                ElmOrderThreadContext.setAll(context);
                runnable.run();
            } finally {
                ElmOrderThreadContext.clear();
            }
        };
    }
}